"""
https://leetcode.cn/problems/maximum-width-of-binary-tree/description/

662. 二叉树最大宽度
中等
相关标签
premium lock icon
相关企业
给你一棵二叉树的根节点 root ，返回树的 最大宽度 。

树的 最大宽度 是所有层中最大的 宽度 。

每一层的 宽度 被定义为该层最左和最右的非空节点（即，两个端点）之间的长度。将这个二叉树视作与满二叉树结构相同，两端点间会出现一些延伸到这一层的 None 节点，这些 None 节点也计入长度。

题目数据保证答案将会在  32 位 带符号整数范围内。

 

示例 1：


输入：root = [1,3,2,5,3,None,9]
输出：4
解释：最大宽度出现在树的第 3 层，宽度为 4 (5,3,None,9) 。
示例 2：


输入：root = [1,3,2,5,None,None,9,6,None,7]
输出：7
解释：最大宽度出现在树的第 4 层，宽度为 7 (6,None,None,None,None,None,7) 。
示例 3：


输入：root = [1,3,2,5]
输出：2
解释：最大宽度出现在树的第 2 层，宽度为 2 (3,2) 。
 

提示：

树中节点的数目范围是 [1, 3000]
-100 <= Node.val <= 100

"""

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
from typing import List, Optional

from tree_node import TreeNode,build_tree_node


class Solution:

    """
    每个节点在每层的位置，最大数量是2^0到2^n,其中n=level-1
    可以分别计算每一层的宽度，之后计算最大宽度

    原方案从数学上讲是可行的，但对空间使用较高。
    考虑换一种方案：改用数字记录空节点数量。（实际使用两个len相同的列表记录）
    在根据上层节点列表，推演下一层节点列表时：
    1. 如果父节点为空，则子级空节点数量*2
    2. 如果父节点不为空，则依据实际情况填充子节点情况
    3. 如果下一次要填充的子节点数量为空，则判断上一个索引是否统计空节点数量。如果是，则在上一个统计数量上+1
    4. 推演层级时，顺便计算宽度（可以解耦做）

    """
    def widthOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        
        def add_new_node(new_node_list:List[TreeNode],new_count_list:List[int],node:TreeNode):
            if node:
                new_node_list.append(node)
                new_count_list.append(1)
            else:
                
                if new_node_list and new_node_list[-1] is None:
                    new_count_list[-1]+=1
                else:
                    new_node_list.append(None)
                    new_count_list.append(1)

        def calc_width(node_list:List[TreeNode],count_list:List[int]):
            start_idx=-1
            end_idx=-1
            # print([(_node.val if _node else None) for _node in node_list])
            # print(count_list)
            # print('----')
            for i in range(len(node_list)):
                if node_list[i]:
                    start_idx=i
                    break
            if start_idx==-1:
                return 0
            
            for i in range(len(node_list)-1,-1,-1):
                if node_list[i]:
                    end_idx=i
                    break
            if start_idx==end_idx:
                return 1
            
            width=0
            for i in range(start_idx,end_idx+1):
                if node_list[i]:
                    width+=1
                else:
                    width+=count_list[i]

            return width

        max_width=0

        def itr_next_level(node_list:List[TreeNode],count_list:List[int]):
            # 此处维护最大宽度
            nonlocal max_width
            max_width=max(max_width,calc_width(node_list,count_list))
            
            # 遍历层级
            new_node_list=[]
            new_count_list=[]
            
            for i,node in enumerate(node_list):
                if not node:
                    # 填充空数量
                    new_count=count_list[i]*2
                    if new_node_list and new_node_list[-1] is None:
                        new_count_list[-1]+=new_count
                    else:
                        new_node_list.append(None)
                        new_count_list.append(new_count)
                else:
                    add_new_node(new_node_list,new_count_list,node.left)
                    add_new_node(new_node_list,new_count_list,node.right)
                    
            if len(new_node_list)==1 and new_node_list[0] is None:
                return

            itr_next_level(new_node_list,new_count_list)

        itr_next_level([root],[1])
        return max_width

        
    
if __name__=='__main__':
    # arr=[0,0,0,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None,None,0,0,None]
    arr=[1,2,3,4,5]
    root=build_tree_node(arr)

    s=Solution()
    res=s.widthOfBinaryTree(root)
    print(res)